#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from __future__ import annotations
import os
import sys
import time
import argparse
from typing import Dict, Any, Tuple

import numpy as np
import astropy.units as u
from astropy.io import fits
from astropy.time import Time
from astropy.wcs import WCS
from astropy.coordinates import SkyCoord
from astroquery.imcce import Skybot


class AsteroidInFITSChecker:
    """
    在一张 FITS 图像中查询给定圆锥范围内是否存在已知小天体（小行星/彗星/卫星），
    并将结果映射到像素坐标，便于落点判定与后续 cutout/标注。

    用法（代码中）:
        checker = AsteroidInFITSChecker(radius_arcmin=50.0, site_code="O44", vmax=21.0)
        result = checker.check("/path/to/your_sciimg.fits")

    结果字典键:
        - 'obs_time': datetime（UTC）
        - 'center_icrs': SkyCoord
        - 'radius_arcmin': float
        - 'table': astropy.table.Table（可能为空）
        - 'x', 'y', 'in_fov'（仅当 add_pixel_coords=True 且有 WCS 时）
    """

    def __init__(self,
                 radius_arcmin: float = 50.0,
                 site_code: str | None = "O44",
                 vmax: float | None = None,
                 timeout: int = 60,
                 retries: int = 2,
                 backoff: float = 1.8,
                 add_pixel_coords: bool = True) -> None:
        """
        :param radius_arcmin: 搜索半径（arcmin）
        :param site_code: MPC 台站码（如丽江高美古 O44）。为 None 时不传入。
        :param vmax: 视星等上限过滤（例如 21.0）。None 表示不过滤。
        :param timeout: SkyBoT 每次请求超时秒数
        :param retries: 失败重试次数
        :param backoff: 指数退避系数
        :param add_pixel_coords: 将结果映射到像素坐标并附加到表格
        """
        self.radius_arcmin = float(radius_arcmin)
        self.site_code = site_code
        self.vmax = vmax
        self.timeout = int(timeout)
        self.retries = int(retries)
        self.backoff = float(backoff)
        self.add_pixel_coords = bool(add_pixel_coords)

        Skybot.TIMEOUT = self.timeout
        Skybot.ROW_LIMIT = -1  # 不限制返回行数

    # -------------------- 公共接口 --------------------

    def check(self, fits_path: str) -> Dict[str, Any]:
        """
        对单个 FITS 执行检索。
        """
        hdr = fits.getheader(fits_path, 0)

        # 观测时间（尽量使用 DATE-OBS）
        date = hdr.get("DATE-OBS") or hdr.get("DATE")
        if date is None:
            raise KeyError(f"{fits_path}: FITS 头缺少 DATE-OBS/DATE")

        obstime_utc = Time(date, scale="utc")
        epoch_tdb_jd = obstime_utc.tdb.jd

        # 中心坐标：优先 TEL_RA/TEL_DEC，其次 CRVAL1/2；再不行用 WCS 图心
        center_icrs, used_wcs = self._get_center_icrs_from_header(hdr)

        # SkyBoT 检索（带重试）
        tbl = self._query_skybot(center_icrs, epoch_tdb_jd)

        # 可能为空
        out: Dict[str, Any] = {
            "obs_time": obstime_utc.to_datetime(),
            "center_icrs": center_icrs,
            "radius_arcmin": self.radius_arcmin,
            "table": tbl
        }
        if len(tbl) == 0:
            return out

        # V 过滤
        if (self.vmax is not None) and ("V" in tbl.colnames):
            tbl = tbl[tbl["V"] <= self.vmax]
            out["table"] = tbl

        # 投影到像素
        if self.add_pixel_coords:
            try:
                xpix, ypix, in_fov = self._project_to_pixels(hdr, tbl)
                tbl["X_SKYBOT"] = xpix
                tbl["Y_SKYBOT"] = ypix
                out.update({"x": xpix, "y": ypix, "in_fov": in_fov})
            except Exception as e:
                # 没有有效 WCS 时允许继续返回天球坐标
                out["projection_error"] = str(e)

        return out

    # -------------------- 具体实现 --------------------

    @staticmethod
    def _get_center_icrs_from_header(hdr) -> Tuple[SkyCoord, bool]:
        ra = hdr.get("TEL_RA") or hdr.get("CRVAL1")
        dec = hdr.get("TEL_DEC") or hdr.get("CRVAL2")

        if (ra is not None) and (dec is not None):
            return SkyCoord(float(ra) * u.deg, float(dec) * u.deg, frame="icrs"), False

        # 退回用 WCS 的图心
        w = WCS(hdr)
        nx, ny = hdr.get("NAXIS1"), hdr.get("NAXIS2")
        if (nx is None) or (ny is None) or (not w.is_celestial):
            raise KeyError("无法获取中心坐标（缺少 TEL_RA/DEC，且无有效 WCS/NAXIS1/2）")
        sky_center = w.pixel_to_world(nx / 2, ny / 2)
        return SkyCoord(sky_center.ra, sky_center.dec, frame="icrs"), True

    def _query_skybot(self, center_icrs: SkyCoord, epoch_tdb_jd: float):
        last_err = None
        for k in range(self.retries + 1):
            try:
                if self.site_code:
                    return Skybot.cone_search(
                        center_icrs,
                        self.radius_arcmin * u.arcmin,
                        epoch_tdb_jd,
                        self.site_code
                    )
                else:
                    return Skybot.cone_search(
                        center_icrs,
                        self.radius_arcmin * u.arcmin,
                        epoch_tdb_jd
                    )
            except Exception as e:
                last_err = e
                if k < self.retries:
                    time.sleep((self.backoff ** k) * 2.0)
                else:
                    raise last_err

    @staticmethod
    def _project_to_pixels(hdr, tbl):
        w = WCS(hdr)
        if not w.is_celestial:
            raise ValueError("WCS 不包含天球坐标轴，无法投影像素。")
        ra_arr = np.array(tbl["RA"], dtype=float)
        dec_arr = np.array(tbl["DEC"], dtype=float)
        sky = SkyCoord(ra_arr * u.deg, dec_arr * u.deg, frame="icrs")
        xpix, ypix = w.world_to_pixel(sky)

        nx, ny = hdr.get("NAXIS1"), hdr.get("NAXIS2")
        if (nx is None) or (ny is None):
            raise ValueError("FITS 缺少 NAXIS1/NAXIS2 以判定视场内外。")
        in_fov = (xpix >= 0) & (xpix < nx) & (ypix >= 0) & (ypix < ny)
        return xpix, ypix, in_fov


# =========================
#        __main__
# =========================
def _iter_fits_files(path: str):
    if os.path.isdir(path):
        for root, _, files in os.walk(path):
            for fn in files:
                if fn.lower().endswith(".fits") and fn.lower().endswith("sciimg.fits"):
                    yield os.path.join(root, fn)
    else:
        yield path


def main():
    ap = argparse.ArgumentParser(
        description="在 FITS 中查询 25 arcmin（可改）范围内的小天体（SkyBoT）并映射像素坐标"
    )
    ap.add_argument("path", help="FITS 文件或目录（目录将递归查找 *sciimg.fits）")
    ap.add_argument("--radius", type=float, default=25.0, help="搜索半径（arcmin），默认 50")
    ap.add_argument("--site", default="O44", help="MPC 台站码（默认丽江 O44），传空串禁用")
    ap.add_argument("--vmax", type=float, default=None, help="按 V 视星等上限过滤（例如 21.0）")
    ap.add_argument("--timeout", type=int, default=60, help="SkyBoT 超时秒数")
    ap.add_argument("--retries", type=int, default=2, help="失败重试次数")
    ap.add_argument("--backoff", type=float, default=1.8, help="指数退避系数")
    ap.add_argument("--no-pix", action="store_true", help="不做像素投影（仅返回天球坐标）")
    ap.add_argument("--csv", default=None, help="将所有结果汇总保存到 CSV")
    args = ap.parse_args()

    site_code = None if (args.site is not None and args.site.strip() == "") else args.site

    checker = AsteroidInFITSChecker(
        radius_arcmin=args.radius,
        site_code=site_code,
        vmax=args.vmax,
        timeout=args.timeout,
        retries=args.retries,
        backoff=args.backoff,
        add_pixel_coords=(not args.no_pix),
    )

    rows = []
    import pandas as pd

    for fpath in _iter_fits_files(args.path):
        try:
            res = checker.check(fpath)
            tbl = res["table"]
            n = len(tbl)
            print(f"\n[{os.path.basename(fpath)}] {res['obs_time']}  center={res['center_icrs'].to_string('hmsdms')}  found={n}")
            if n > 0:
                # 打印前 5 个目标
                keep = [c for c in ["Number", "Name", "Type", "RA", "DEC", "V", "X_SKYBOT", "Y_SKYBOT"] if c in tbl.colnames]
                print(tbl[keep])

                # 收集到 rows 里
                df = tbl.to_pandas()
                df.insert(0, "fitsname", os.path.basename(fpath))
                df.insert(1, "obs_time", pd.Timestamp(res["obs_time"]))
                rows.append(df)
        except Exception as e:
            print(f"[ERROR] {fpath}: {e}", file=sys.stderr)

    if rows:
        big = pd.concat(rows, ignore_index=True)
        out_csv = args.csv or "asteroid_query_results.csv"
        big.to_csv(out_csv, index=False)
        print(f"\nSaved CSV: {out_csv}")
    else:
        print("\nNo objects found or no valid results.")


if __name__ == "__main__":
    main()



# 这很常见——用本地 MPCORB.DAT 推进通常会“比 SkyBoT 多”。原因大致有几类：

# 为啥会“更多”

# 对象覆盖面不同

# 你的离线库 = MPCORB.DAT 全量（含很多轨道很差、只观测过一次/短弧的临时目标）。

# SkyBoT 使用 IMCCE 的交叉整合库，并清洗/去重/合并别名；也可能暂缓返回轨道过差的候选。

# 历元 & 推进模型差异

# PyEphem 按 MPCORB 的近拱点轨道在单一历元推进；

# SkyBoT 采用其内部星历（批量数值积分/更频繁更新），对高不确定度目标更保守。
# 结果：有些 MPCORB 计算到的天区命中，SkyBoT 认为不够确定/不在场。

# 观测地点/视差/框选方式

# 你用的是本地站点 topocentric（经纬高已设）；

# SkyBoT 若没传台站码，可能按地心；或者你传了也可能与本地有微小差别。

# 圆锥半径：你/我代码用 ICRS 圆锥（50′）；SkyBoT 也是，但若你中心坐标来自不同头关键字（TEL_RA/CRVAL/WCS），会有中心偏差导致边缘对象计入/排除不同。

# 亮度阈值 & 列表截断

# 你的离线结果默认不过滤亮度；

# 你若在 SkyBoT 打过 --vmax，就会少；反之 SkyBoT 也可能对极暗目标估计不同。

# SkyBoT 在极密区（靠黄道/银河）偶有响应超时或返回截断（虽已设 ROW_LIMIT=-1，服务器仍可能限流）。